Visualisation

Note

This example is available as a jupyter notebook here.

import ring
from ring import exp
import mediapy as media
import jax
sys_str = """
<x_xy>
<worldbody>
<geom dim="0.1" type="xyz"></geom>
<body joint="free" name="seg" pos="0 0 .5">
<geom color="dustin_exp_blue" dim="0.15 0.075 0.05" mass="0.2" pos="0.03 0 0" type="box"></geom>
<body joint="frozen" name="imu" pos="0.0 0.0 0.03">
<geom color="dustin_exp_orange" dim="0.05 0.03 0.02" mass="0.1" type="box"></geom>
</body>
</body>
</worldbody>
</x_xy>
"""
sys = ring.System.create(sys_str)
(X, y), (key, q, x, _) = ring.RCMG(sys, keep_output_extras=True).to_list()[0]
eager data generation: 1it [00:01,  1.58s/it]

media.show_video(sys.render(x, width=640, height=480, camera="target", render_every_nth=4), fps=25)
Rendering frames..: 100%|██████████| 1500/1500 [00:08<00:00, 167.58it/s]

exp_data = exp.load_data(exp_id="S_06", motion_start="fast")
exp_data.keys()
dict_keys(['seg1', 'seg2', 'seg3', 'seg4', 'seg5'])
exp_data["seg1"].keys()
dict_keys(['imu_flex', 'imu_rigid', 'marker1', 'marker2', 'marker3', 'marker4', 'quat'])
segment = "seg2"
omc_data_sys = {
    "seg": {
        "pos": exp_data[segment]["marker1"],
        "quat": exp_data[segment]["quat"],
    },
    "imu": {
        "quat": exp_data[segment]["quat"],
    }
}
omc_data_sys
{'seg': {'pos': Array([[ 0.26832035,  1.1925832 , -0.06244465],
         [ 0.268319  ,  1.1925788 , -0.06244366],
         [ 0.26831627,  1.1925731 , -0.06244285],
         ...,
         [ 0.19899347,  1.2143577 , -0.06264979],
         [ 0.19898805,  1.2143621 , -0.06264362],
         [ 0.19897974,  1.214368  , -0.06263389]], dtype=float32),
  'quat': Array([[ 0.95449424,  0.10144146, -0.01664882, -0.27995223],
         [ 0.95449567,  0.10146226, -0.01665274, -0.27993944],
         [ 0.95449716,  0.10148306, -0.01665665, -0.27992663],
         ...,
         [ 0.9600311 , -0.01755334,  0.02215116, -0.2784628 ],
         [ 0.9600333 , -0.01756094,  0.02225687, -0.27844635],
         [ 0.96003544, -0.01756854,  0.02236257, -0.27842987]],      dtype=float32)},
 'imu': {'quat': Array([[ 0.95449424,  0.10144146, -0.01664882, -0.27995223],
         [ 0.95449567,  0.10146226, -0.01665274, -0.27993944],
         [ 0.95449716,  0.10148306, -0.01665665, -0.27992663],
         ...,
         [ 0.9600311 , -0.01755334,  0.02215116, -0.2784628 ],
         [ 0.9600333 , -0.01756094,  0.02225687, -0.27844635],
         [ 0.96003544, -0.01756854,  0.02236257, -0.27842987]],      dtype=float32)}}
x = ring.sim2real.xs_from_raw(sys, omc_data_sys)

# vectorize this function over time
@jax.vmap
def update_position_vector_of_imu(x):
    state = ring.State.create(sys, x=x)
    # populate minimal coordinates `state.q` from maximal coordinates `state.x`
    state = ring.algorithms.inverse_kinematics(sys, state)
    # re-calculate maximal coordiantes `state.x` from minimal coordinates `state.q`
    # this uses the position vector specified in the system (and so the xml file)
    # to produce an offset between IMu and segment geom box
    _, state = ring.algorithms.forward_kinematics(sys, state)
    return state.x

x = update_position_vector_of_imu(x)
media.show_video(sys.render(x, width=640, height=480, camera="target", render_every_nth=4), fps=25)
Rendering frames..: 100%|██████████| 1075/1075 [00:07<00:00, 134.95it/s]